/**
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *
 * @flow
 * @format
 * @oncall react_native
 */

import type {Writable} from 'stream';

import {startProfiling, stopProfilingAndWrite} from './profiling';
import JSONStream from './third-party/JSONStream';
import {Console} from 'console';
import duplexer from 'duplexer';
import fs from 'fs';
import invariant from 'invariant';

export type Command = (
  argv: Array<string>,
  structuredArgs: mixed,
  console: Console,
) => Promise<void> | void;
export type Commands = {[key: string]: Command, ...};

type Message<Type: string, Data> = Data & {
  id: number,
  type: Type,
  ...
};

type HandshakeMessage = Message<
  'handshake',
  {
    protocol_version: '0',
    capabilities: [],
    ...
  },
>;

type CommandMessage = Message<
  'command',
  {
    args_path: string,
    stdout_path: string,
    stderr_path: string,
    ...
  },
>;

type HandshakeReponse = Message<
  'handshake',
  {
    protocol_version: '0',
    capabilities: [],
    ...
  },
>;

type CommandResponse = Message<'result', {exit_code: 0, ...}>;

type ErrorResponse = Message<'error', {exit_code: number, ...}>;

type IncomingMessage = HandshakeMessage | CommandMessage;
type Response = HandshakeReponse | CommandResponse | ErrorResponse;
type RespondFn = (response: Response) => void;

type JSONReaderDataHandler = IncomingMessage => mixed;
type JSONReaderEndHandler = () => mixed;

type JSONReaderDataListener = ('data', JSONReaderDataHandler) => JSONReader;
type JSONReaderEndListener = ('end', JSONReaderEndHandler) => JSONReader;
type JSONReaderRootEndListener = (
  'root_end',
  JSONReaderEndHandler,
) => JSONReader;
type JSONReaderListener = JSONReaderDataListener &
  JSONReaderEndListener &
  JSONReaderRootEndListener;

type JSONReader = {
  on: JSONReaderListener,
  removeListener: JSONReaderListener,
  ...
};

type JSONWriter = {
  write(object: Response): void,
  end(object?: Response): void,
  ...
};

function buckWorker(commands: Commands): any {
  const reader: JSONReader = JSONStream.parse('*');
  const writer: JSONWriter = JSONStream.stringify();

  function handleHandshake(message: IncomingMessage): void {
    const response = handshakeResponse(message);

    if (response.type === 'handshake') {
      if (JS_WORKER_TOOL_CPU_PROFILE) {
        // $FlowFixMe[unused-promise]
        startProfiling().then(() => writer.write(response));
      } else {
        writer.write(response);
      }
      reader.removeListener('data', handleHandshake).on('data', handleCommand);
    } else {
      writer.write(response);
    }
  }

  function handleCommand(message: IncomingMessage): void {
    const {id} = message;

    if (message.type !== 'command') {
      writer.write(unknownMessage(id));
      return;
    }

    if (!message.args_path || !message.stdout_path || !message.stderr_path) {
      writer.write(invalidMessage(id));
      return;
    }

    let responded: boolean = false;
    let stdout, stderr;

    try {
      stdout = fs.createWriteStream(message.stdout_path);
      stderr = fs.createWriteStream(message.stderr_path);
    } catch (e) {
      respond(invalidMessage(id));
      return;
    }

    readArgsAndExecCommand(message, commands, stdout, stderr, respond);

    function respond(response: Response) {
      // 'used for lazy `.stack` access'
      invariant(!responded, `Already responded to message id ${id}.`);
      responded = true;

      void Promise.all(
        [stdout, stderr]
          .filter(Boolean)
          .map(stream => new Promise(resolve => stream.end(resolve))),
      ).then(() => {
        writer.write(response);
      });
    }
  }

  let ended = false;
  function end() {
    if (ended) {
      return;
    }
    ended = true;
    // $FlowFixMe[unused-promise]
    stopProfilingAndWrite(JS_WORKER_TOOL_NAME).then(() => writer.end());
  }
  reader.on('data', handleHandshake);
  reader.on('end', end);
  reader.on('root_end', end);
  return duplexer(reader, writer);
}

function handshakeResponse(message: IncomingMessage) {
  if (message.type !== 'handshake') {
    return unknownMessage(message.id);
  }

  if (message.protocol_version !== '0') {
    return invalidMessage(message.id);
  }

  return {
    id: message.id,
    type: 'handshake' as const,
    protocol_version: '0' as const,
    capabilities: [] as [],
  };
}

function readArgsAndExecCommand(
  message: CommandMessage,
  commands: Commands,
  stdout: Writable,
  stderr: Writable,
  respond: RespondFn,
) {
  const {id} = message;

  fs.readFile(message.args_path, 'utf8', (readError, argsString) => {
    if (readError) {
      respond(invalidMessage(id));
      return;
    }

    let commandName;
    let args: Array<string> = [];
    let structuredArgs = null;

    // If it starts with a left brace, we assume it's JSON-encoded. This works
    // because the non-JSON encoding always starts the string with the
    // command name, thus a letter.
    if (argsString[0] === '{') {
      ({command: commandName, ...structuredArgs} = JSON.parse(argsString));
    } else {
      // FIXME: if there are files names with escaped
      // whitespace, this will not work.
      [commandName, ...args] = argsString.split(/\s+/);
    }

    if (commands.hasOwnProperty(commandName)) {
      const command = commands[commandName];
      const commandSpecificConsole = new Console(stdout, stderr);
      // $FlowFixMe[unused-promise]
      execCommand(
        command,
        commandName,
        argsString,
        args,
        structuredArgs,
        commandSpecificConsole,
        respond,
        id,
      );
    } else {
      stderr.write(
        `This worker does not have a command named \`${commandName}\`. ` +
          `Available commands are: ${Object.keys(commands).join(', ')}`,
      );
      respond(invalidMessage(id));
    }
  });
}

const {
  JS_WORKER_TOOL_DEBUG_RE,
  JS_WORKER_TOOL_CPU_PROFILE,
  JS_WORKER_TOOL_NAME,
} = process.env;
const DEBUG_RE = JS_WORKER_TOOL_DEBUG_RE
  ? new RegExp(JS_WORKER_TOOL_DEBUG_RE)
  : null;

async function execCommand(
  command: Command,
  commandName: string,
  argsString: string,
  args: Array<string>,
  structuredArgs: mixed,
  commandSpecificConsole: Console,
  respond: RespondFn,
  messageId: number,
) {
  let makeResponse: (id: number) => Response = success;
  try {
    if (shouldDebugCommand(argsString)) {
      throw new Error(
        `Stopping for debugging. Command '${commandName} ...' matched by the 'JS_WORKER_TOOL_DEBUG_RE' environment variable`,
      );
    }
    await command(args.slice(), structuredArgs, commandSpecificConsole);
  } catch (e) {
    commandSpecificConsole.error(e.stack);
    makeResponse = commandError;
  }

  respond(makeResponse(messageId));
}

function shouldDebugCommand(argsString: string) {
  return DEBUG_RE && DEBUG_RE.test(argsString);
}

const error = (id: number, exitCode: number) => ({
  type: 'error' as const,
  id,
  exit_code: exitCode,
});
const unknownMessage = (id: number) => error(id, 1);
const invalidMessage = (id: number) => error(id, 2);
const commandError = (id: number) => error(id, 3);
const success = (id: number) => ({
  type: 'result' as const,
  id,
  exit_code: 0 as const,
});

export {buckWorker};
